import asyncio
import json
import os.path

import websockets
from playwright.async_api import async_playwright, Browser, BrowserContext


class PlaywrightManager:

    def __init__(self):
        self.browser: Browser = None
        self.context: BrowserContext = None

    async def start(self):
        playwright = await async_playwright().__aenter__()
        self.browser = await playwright.chromium.launch(headless=False)
        if not os.path.exists('cookies.json'):
            with open('cookie.json', 'w') as f:
                json.dump({'1': 2}, f)
        self.context = await self.browser.new_context(storage_state="cookie.json")

    async def stop(self):
        await self.context.close()
        await self.browser.close()

    async def Send(self, secuid: str, msg: str):
        if not self.context:
            await self.start()
        page = await self.context.new_page()
        url = "https://www.douyin.com/user/" + secuid
        await page.goto(url, timeout=60000)
        await page.wait_for_selector('//p[@class="HkAZSS6r" and contains(text(),"私信")]', timeout=50000)
        await asyncio.sleep(5)
        await page.click('//div[@class="hLIm2dFu"]//span[contains(text(),"私信")]')
        await page.wait_for_selector('//div[@class="public-DraftStyleDefault-block public-DraftStyleDefault-ltr"]')
        await asyncio.sleep(3)
        input_element = await page.query_selector(
            '//div[@class="public-DraftStyleDefault-block public-DraftStyleDefault-ltr"]')
        await input_element.fill(msg)
        await page.click('//span[@class="sCp7KhBv EWT1TDgs e2e-send-msg-btn"]')
        await page.close()

    async def check_login(self):
        if not self.context:
            await self.start()
        page = await self.context.new_page()
        await page.goto("http://www.douyin.com")
        try:
            element = await page.wait_for_selector(
                '//img[@class="web-login-scan-code__content__qrcode-wrapper__qrcode"]')
            result = False
        except:
            result = True

        if result:
            print("已登录")
            await page.close()
            return False
        else:
            # qrcode = await page.query_selector(
            #     '//img[@class="web-login-scan-code__content__qrcode-wrapper__qrcode"]')
            # await qrcode.screenshot(path='../templates/static/qrcode.png')
            # await page.wait_for_selector('//p[contains(text(), "登录成功")]')
            # await asyncio.sleep(2)
            input("登录成功后在这里按下回车......")
            storage_state = await self.context.storage_state()
            with open('cookie.json', 'w') as f:
                json.dump(storage_state, f)

            print("登录成功")
            await page.close()
            return True


playwright_manager = PlaywrightManager()


async def websocket_client():
    uri = "ws://47.113.223.17:80/ws"  # WebSocket 服务器地址
    async with websockets.connect(uri) as websocket:
        print("Connected to WebSocket server")
        while True:
            try:
                message = await websocket.recv()
                await asyncio.create_task(playwright_manager.Send(**json.loads(message)))
            except websockets.ConnectionClosed:
                print("Connection closed")
                break
            except Exception as e:
                print(e)
                continue


async def main():
    await playwright_manager.start()
    await playwright_manager.check_login()
    await websocket_client()


if __name__ == '__main__':
    asyncio.run(main())
